import math
import numpy as np
import gurobipy as grb
from utils.base import ArchBase


class Group_LP(ArchBase):
    """
    所有Rack间均有直连链路，同一时刻满足分组后的下行需求。
    """
    def __init__(self, **kwargs):
        super().__init__(**kwargs)


    def calc_min_provision(self, demands):

        m = grb.Model(self.__class__.__name__)
        m.Params.LogToConsole = 0

        uplink_in_rack = m.addVar(lb = 0, vtype = grb.GRB.CONTINUOUS, name = 'U')
        rack_topo = m.addVars(self.rack_num, self.rack_num, vtype = grb.GRB.CONTINUOUS, name = 'x')
        
        name_w_ikj = [
            f'w_{i}_{k}_{j}'
            for i in range(1, self.rack_num + 1)
            for k in range(self.rack_num + 1)
            for j in range(1, self.rack_num + 1)
            if i != k and i != j and j != k
        ]
        w = m.addVars(name_w_ikj, lb = 0, vtype = grb.GRB.CONTINUOUS, name = 'w')

        # summation(x_ij) <= U
        m.addConstrs(
            grb.quicksum(
                rack_topo[i, j] for j in range(self.rack_num)
            ) <= uplink_in_rack
            for i in range(self.rack_num)
        )
        m.addConstrs(
            grb.quicksum(
                rack_topo[j, i] for j in range(self.rack_num)
            ) <= uplink_in_rack
            for i in range(self.rack_num)
        )
        
        # 控制对角位置元素相等
        m.addConstrs(
            rack_topo[i, j] == rack_topo[j, i]
            for j in range(self.rack_num)
            for i in range(self.rack_num)
            if i != j
        )

        # summation(w_ikj) = 1
        m.addConstrs(
            (grb.quicksum(
                w[f'w_{i}_{k}_{j}'] for k in range(0, self.rack_num + 1)
                if k != i and k != j
            ) == 1
            for i in range(1, self.rack_num + 1)
            for j in range(1, self.rack_num + 1)
            if i != j),
            name = 'sumOneConstrs'
        )

        # m.addConstrs(
        #     grb.quicksum(
        #         demands * w[f'w_{k}_{i}_{j}']
        #         + demands * w[f'w_{i}_{j}_{k}']
        #         if k != 0 else demands * w[f'w_{i}_{0}_{j}']
        #         for k in range(0, self.rack_num + 1)
        #         if k != i and k != j
        #     ) <= rack_topo[i - 1, j - 1]
        #     for i in range(1, self.rack_num + 1)
        #     for j in range(1, self.rack_num + 1)
        #     if i != j
        # )
        
        # Grouping
        mask = 1
        # self.rack_num = 8  # debug
        for _ in range(0, int(math.log(self.rack_num, 2))):
            demands_TM = np.zeros((self.rack_num, self.rack_num))
            for pair in range(0, self.rack_num):
                # print(pair, (pair ^ mask))
                # 此处相当于实际Rack编号-1
                rack_i = pair
                rack_j = (pair ^ mask)
                demands_TM[rack_i][rack_j] = demands

            m.addConstrs(
                grb.quicksum(
                    demands_TM[k - 1][j - 1] * w[f'w_{k}_{i}_{j}']
                    + demands_TM[i - 1][k - 1] * w[f'w_{i}_{j}_{k}']
                    if k != 0 else demands_TM[i - 1][j - 1] * w[f'w_{i}_{0}_{j}']
                    for k in range(0, self.rack_num + 1)
                    if k != i and k != j
                ) <= rack_topo[i - 1, j - 1]
                for i in range(1, self.rack_num + 1)
                for j in range(1, self.rack_num + 1)
                if i != j
            )
            
            # print()
            mask = mask << 1


        m.setObjective(uplink_in_rack, grb.GRB.MINIMIZE)
        m.optimize()
        # m.write('debug.lp')
        if m.status == grb.GRB.Status.OPTIMAL:
            print(self.__class__.__name__, m.objVal)

            rack_topo_solution = m.getAttr('X', rack_topo)
            print(rack_topo_solution)
            w_solution = m.getAttr('X', w)
            print(w_solution)
            for k, v in dict(w_solution).items():
                if v > 0:
                    print(k, v)
            self.uplink_in_rank = m.objVal
        else:
            print(self.__class__.__name__ + 'is no solution')
            
        return self.uplink_in_rank

